#!/usr/bin/env python3
import math
import rclpy
from rclpy.node import Node
from tf2_ros import TransformBroadcaster, TransformListener, Buffer
from geometry_msgs.msg import TransformStamped
from nav_msgs.msg import Odometry


class Tag2Chassis(Node):
    def __init__(self):
        super().__init__("tag2chassis")
        self.get_logger().info("tag2chassis node started")
        self.declare_parameter("frame_prefix", "")
        self.declare_parameter("target_prefix", "")
        self.frame_prefix: str = self.get_parameter("frame_prefix").value
        self.target_prefix: str = self.get_parameter("target_prefix").value
        if self.frame_prefix[0] == "/":
            self.frame_prefix = self.frame_prefix[1:]
        if self.target_prefix[0] == "/":
            self.target_prefix = self.target_prefix[1:]

        self.tf_broadcaster = TransformBroadcaster(self)
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, self)

        self.odomPub = self.create_publisher(Odometry, "tag_odom", 10)

        self.create_timer(0.1, self.publish_tag_chassis_tf)

    def publish_tag_chassis_tf(self):
        try:
            # 获取 chassis 到 tag 的 TF
            tag_to_chassis = self.tf_buffer.lookup_transform(
                self.target_prefix + "/stag",
                self.target_prefix + "/chassis",
                rclpy.time.Time(),
            )

            # 创建并发布 tag_chassis 的 TF
            tag_chassis_tf = TransformStamped()
            tag_chassis_tf.header.stamp = self.get_clock().now().to_msg()
            tag_chassis_tf.header.frame_id = self.frame_prefix + "/tag"
            tag_chassis_tf.child_frame_id = self.frame_prefix + "/tag_chassis"
            tag_chassis_tf.transform = tag_to_chassis.transform

            self.tf_broadcaster.sendTransform(tag_chassis_tf)

            # 发布 Odometry 消息（可选）
            chassis2map = self.tf_buffer.lookup_transform(
                "map",
                self.frame_prefix + "/tag_chassis",
                rclpy.time.Time(),
            )
            odom = Odometry()
            odom.header.stamp = self.get_clock().now().to_msg()
            odom.header.frame_id = "map"
            odom.child_frame_id = self.frame_prefix + "/tag_chassis"
            odom.pose.pose.position.x = chassis2map.transform.translation.x
            odom.pose.pose.position.y = chassis2map.transform.translation.y
            odom.pose.pose.position.z = chassis2map.transform.translation.z
            odom.pose.pose.orientation = chassis2map.transform.rotation
            self.odomPub.publish(odom)

        except Exception as e:
            #self.get_logger().warn(f"Failed to lookup transform: {e}")
            pass


def main(args=None):
    rclpy.init(args=args)
    node = Tag2Chassis()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()


if __name__ == "__main__":
    main()
